// 文件路径: dbms/cluster/ShardingManager.java
package com.dbms.cluster;

import com.dbms.parser.ast.ExpressionNode;
import com.dbms.parser.ast.StatementNode;
import com.dbms.parser.ast.DeleteNode;
import com.dbms.parser.ast.InsertNode;
import com.dbms.parser.ast.SelectNode;
import com.dbms.parser.ast.UpdateNode;

import java.util.List;
import java.util.ArrayList;
import java.util.Collections;

/**
 * 分片管理器 (Sharding Manager)。
 * 负责实现数据分片的核心逻辑。
 * 主要功能包括：
 * 1. 根据分片键的值，计算出它应该被路由到哪个物理节点。
 * 2. 从SQL语句的抽象语法树（AST）中提取出分片键的值。
 */
public class ShardingManager {

    // 存储一个有序的节点列表。排序是为了保证所有节点上的路由算法都是确定和一致的。
    private final List<NodeInfo> nodes;

    /**
     * ShardingManager 的构造函数。
     */
    public ShardingManager() {
        // 从 ClusterManager 获取集群中所有节点的列表。
        List<NodeInfo> sortedNodes = new ArrayList<>(ClusterManager.getInstance().getClusterNodes().values());
        // 【关键】对节点列表进行排序（例如，按节点ID的字典序）。
        // 这一步至关重要，因为它确保了在集群中的任何一个节点上，对于同一个分片键，
        // 计算出的目标节点总是一样的，保证了路由的一致性。
        Collections.sort(sortedNodes, (n1, n2) -> n1.getNodeId().compareTo(n2.getNodeId()));
        this.nodes = sortedNodes;
    }

    /**
     * 根据分片键的值，计算它应该被存储或查询的节点。
     * 这是一个典型的哈希取模路由算法。
     * @param shardingKey 分片键的值 (这里假设它是一个整数或可以计算哈希值的对象)。
     * @return 应该处理该键的目标节点的 NodeInfo。
     */
    public NodeInfo getNodeForKey(Object shardingKey) {
        if (nodes.isEmpty()) {
            throw new IllegalStateException("集群中没有可用的节点。");
        }
        // 1. 计算分片键的哈希值。
        int hashCode = shardingKey.hashCode();
        // 2. 用哈希值的绝对值对节点总数取模，得到目标节点在有序列表中的索引。
        int index = Math.abs(hashCode % nodes.size());
        // 3. 返回计算出的目标节点。
        return nodes.get(index);
    }

    /**
     * 从SQL语句的抽象语法树（AST）中提取分片键的值。
     * 这是一个简化的实现，它基于一些预设的规则来寻找分片键。
     * @param stmt SQL语句解析后生成的AST节点。
     * @return 提取出的分片键的值。如果找不到（例如，一个不带WHERE的全表扫描），则返回 null。
     */
    public static Object extractShardingKey(StatementNode stmt) {
        if (stmt instanceof InsertNode) {
            InsertNode node = (InsertNode) stmt;
            // 假设：对于INSERT语句，第一个插入的值就是主键/分片键。
            if (!node.getValues().isEmpty()) {
                // 获取第一个值，它可能带有引号，需要移除。
                String value = node.getValues().get(0).replace("'", "");
                try {
                    // 尝试将其转换为整数，因为分片键通常是数字ID。
                    return Integer.parseInt(value);
                } catch (NumberFormatException e) {
                    // 如果转换失败，则按原样返回字符串（支持字符串类型的主键）。
                    return value;
                }
            }
        } else if (stmt instanceof SelectNode || stmt instanceof DeleteNode || stmt instanceof UpdateNode) {
            // 对于 SELECT, DELETE, UPDATE，我们从 WHERE 子句中寻找分片键。
            ExpressionNode whereClause = null;
            if (stmt instanceof SelectNode) whereClause = ((SelectNode) stmt).getWhereClause();
            else if (stmt instanceof DeleteNode) whereClause = ((DeleteNode) stmt).getWhereClause();
            else if (stmt instanceof UpdateNode) whereClause = ((UpdateNode) stmt).getWhereClause();

            // 假设：WHERE子句是一个简单的等式，如 'pk_column = value'。
            if (whereClause != null && whereClause.getOperator().equals("=")) {
                // 这是一个非常简化的假设，我们直接返回等号右边的操作数作为分片键的值。
                // 一个更完整的实现需要检查左操作数是否确实是预定义的分片键列。
                return whereClause.getRightOperand();
            }
        }
        // 如果SQL语句没有明确指向单个分片键（例如 SELECT * FROM table;），则返回null。
        // 这种查询通常需要广播到所有节点执行。
        return null;
    }
}